import cv2
import imutils
import os
from utils.visualize import vis, label2chinese
from detector_yolov7 import Detector
from chinese_field import DetectClasses


def detect_photo(file_pic, weight='', save_image_path=''):
    print(file_pic)
    detect_pic = cv2.imread(file_pic)
    (pic_height, pic_width, _) = detect_pic.shape
    resize_h = pic_height
    if int(resize_h) != int(pic_height):
        im = imutils.resize(detect_pic, height=resize_h)
    else:
        im = detect_pic
    print('HEIGHT:', pic_height)
    print('WIDTH:', pic_width)
    print(save_image_path)

    detector = Detector(weight=weight, conf_thres=0.7, iou_thres=0.5)
    _, detect_info = detector.detect(im)
    if detect_info['box_nums'] > 0:
        img = detect_info['raw_img']
        for order in range(detect_info['box_nums']):
            boxes = detect_info['boxes'][order]
            x0 = int(boxes[0])
            y0 = int(boxes[1])
            x1 = int(boxes[2])
            y1 = int(boxes[3])
            cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), 2)
            class_ids = detect_info['class_ids'][order]
            class_name = detector.names[int(class_ids)]
            img = label2chinese(img,
                                x0, y0, x1, y1,
                                DetectClasses[class_name])
        detect_info['visual'] = img
    else:
        detect_info['visual'] = detect_info['raw_img']
    if save_image_path != '':
        cv2.imwrite(save_image_path, detect_info['visual'])


if __name__ == '__main__':
    photo_direc = '/home/kaijia/algo-env/datasets/214'
    photo_list = os.listdir(photo_direc)
    photo_num = len(photo_list)
    photo_count = 0
    for photo in photo_list:
        photo_count += 1
        photo_origin = os.path.join(photo_direc, photo)
        # photo_save = os.path.join(photo_direc, 'demo_' + photo)
        photo_save = 'demo_' + photo
        detect_photo(photo_origin, weight='weights/yolov7_hq_safeguard.pt',
                     save_image_path=os.path.join("./runs/output", photo_save))
        print('Processing: ', photo_num, '/', photo_count)
